//
// Created by liuhuan on 2021/5/23.
//

#include "Hgd.h"
#include "math.h"

using namespace std;

uint Hgd::rhyper(uint kk, uint nn1, uint nn2, Coins coins) {
    if (kk > 10) {
        return Hgd::hypergeometricHrua(coins, nn1, nn2, kk);
    } else {
        return Hgd::hypergeometricHyp(coins, nn1, nn2, kk);
    }
}

uint Hgd::hypergeometricHyp(Coins coins, uint good, uint bad, uint sample) {
    double d1 = bad + good - sample;
    double d2 = fmin(bad, good);
    double Y = d2;
    double K = sample;
    while (Y > 0.0) {
        double U = coins.draw();
        Y -= int(floor(U + Y/(d1 + K)));
        K -= 1;
        if (K == 0) {
            break;
        }
    }
    uint Z = int(d2 - Y);
    if (good > bad) {
        Z = sample - Z;
    }
    return Z;
}

uint Hgd::hypergeometricHrua(Coins coins, uint good, uint bad, uint sample) {
    double D1 = 1.7155277699214135;
    double D2 = 0.8989161620588988;

    uint minGoodBad = fmin(good, bad);
    uint popSize = good + bad;
    uint maxGoodBad = fmax(good, bad);
    uint m = fmin(sample, popSize - sample);

    double d4 = (float)minGoodBad / popSize;
    double d5 = 1.0 - d4;
    double d6 = m*d4 + 0.5;
    double d7 = sqrt((popSize - m) * sample * d4 * d5 / (popSize-1) + 0.5);
    double d8 = D1*d7 + D2;
    uint d9 = uint(floor(float(m+1) * (minGoodBad+1) / (popSize+2)));
    double d10 = Hgd::loggam(d9+1) + Hgd::loggam(minGoodBad-d9+1) + Hgd::loggam(m-d9+1) + Hgd::loggam(minGoodBad-m+d9+1);
    double d11 = fmin(fmin(m, minGoodBad) + 1.0, floor(d6 + 16 * d7));

    double Z;
    while (true) {
        double X = coins.draw();
        double Y = coins.draw();
        double W = d6 + d8 * (Y -0.5) / X;

        if (W < 0.0 or W >= d11) {
            continue;
        }

        Z = int(floor(W));
        double T = d10 - (Hgd::loggam(Z+1) + Hgd::loggam(minGoodBad-Z+1) + Hgd::loggam(m-Z+1) + Hgd::loggam(maxGoodBad-m+Z+1));

        if ((X*(4.0-X)-3.0) <= T) {
            break;
        }

        if (X*(X-T) >= 1) {
            continue;
        }

        if (2.0 * log(X) <= T) {
            break;
        }
    }

    if (good > bad) {
        Z = m - Z;
    }

    if (m < sample) {
        Z = good - Z;
    }

    return Z;
}

double Hgd::loggam(double x) {
    double a[] = {8.333333333333333e-02, -2.777777777777778e-03,
                  7.936507936507937e-04, -5.952380952380952e-04,
                  8.417508417508418e-04, -1.917526917526918e-03,
                  6.410256410256410e-03, -2.955065359477124e-02,
                  1.796443723688307e-01, -1.39243221690590e+00};
    double x0 = x;
    int n = 0;
    if (x == 1.0 or x == 2.0) {
        return 0.0;
    } else if (x <= 7.0) {
        n = int(7 - x);
        x0 = x + n;
    }
    double x2 = 1.0 / (x0 * x0);
    double gl0 = a[9];
    for (int k = 8; k > -1; k--) {
        gl0 *= x2;
        gl0 += a[k];
    }
    double gl = gl0/x0 + 0.5 * log(2*PI) + (x0-0.5) * log(x0) - x0;
    if (x <= 7.0) {
        for (int k = 1; k < n+1; k++) {
            gl -= log(x0-1.0);
            x0 -= 1.0;
        }
    }
    return gl;
}
